import os, signal, subprocess, sys, time
import cPickle as pickle
import pickle
import multiprocessing

from collections import defaultdict
from StringIO import StringIO


USE_PGRP = False
VERBOSE = False

def all_subprocesses(root_pid):
    listing = subprocess.Popen(['ps', '-o', 'pid,ppid'], stdout=subprocess.PIPE).communicate()[0].split('\n')[1:]
    children = defaultdict(list)
    for line in listing:
        line = line.strip()
        if line:
            pid, ppid = [int(s.strip()) for s in line.split(' ') if s]
            children[ppid].append(pid)
    all = []
    def add_children(ppid):
        for pid in children[ppid]:
            all.append(pid)
            add_children(pid)
    add_children(root_pid)
    return all

def kill_all(pids, sig):
    for pid in pids:
        try:
            os.kill(pid, sig)
        except:
            pass

def really_kill(p):
    #        os.system("ps -fTj")
    #        os.system("./pstree -p %s" % os.getpid())
    #        os.system("kill -9 -%s" % p.pid)
    #        os.system("kill -9 %s" % p.pid)
    #        os.killpg(p.pid, signal.SIGSTOP)
    if p.is_alive() and VERBOSE:
        print "killing", p.pid
    if USE_PGRP:
        p.terminate()
        p.join(1)
        os.killpg(p.pid, signal.SIGTERM)
        time.sleep(.1)
        os.killpg(p.pid, signal.SIGKILL)
        p.join(1)
    elif p.is_alive():
        subprocesses = all_subprocesses(p.pid)
        if VERBOSE:
            print "subprocesses", subprocesses
        p.terminate()
        p.join(1)
        if subprocesses:
            kill_all(reversed(subprocesses), signal.SIGTERM)
            time.sleep(.1)
            kill_all(reversed(subprocesses), signal.SIGKILL)
            p.join(1)
    if p.is_alive():
        os.kill(p.pid, signal.SIGKILL)
        p.join(1)

def fork_and_run(func, timeout=None, streaming=False):
#    parent_pipe, child_pipe = multiprocessing.Pipe(False)
    result_file = os.tmpfile() # pipes block
    if not streaming:
        output = os.tmpfile()
    else:
        output = StringIO()
    def do_it():
        if USE_PGRP:
            os.setpgrp()
        if not streaming:
            os.dup2(output.fileno(), sys.stdout.fileno())
            os.dup2(output.fileno(), sys.stderr.fileno())
        result = func()
        try:
            pickle.dump(result, result_file)
            result_file.flush()
        except:
            result_file.seek(0)
            raise
    p = multiprocessing.Process(target=do_it)
    try:
        p.start()
        p.join(timeout)
        if result_file.tell() == 0:
            result = None
        else:
            result_file.seek(0)
            result = pickle.load(result_file)
        output.seek(0)
        return p.is_alive(), p.exitcode, result, output.read()
    finally:
        really_kill(p)
